import re
from copy import copy
from typing import Dict, List

CHARS = list(range(65, 91))


class Enigma(object):
    def __init__(self, wheels: List[List], wheels_init: List, wheels_carry: List, single_table: Dict = None):
        self.__wheels_carry = wheels_carry
        self.__single_table = single_table
        self.__reflector = list(zip(CHARS, CHARS[int(len(CHARS) / 2):] + CHARS[:int(len(CHARS) / 2)]))  # 反射板
        self.__wheels = []
        for i in range(len(wheels)):
            wheel_ready = copy(wheels[i])
            while wheel_ready[0][0] != wheels_init[i]:
                wheel_ready = self.__step(wheel_ready)
            self.__wheels.append(wheel_ready)

    @property
    def current_wheel_state(self):
        wheel_state = []
        for wheel in self.__wheels:
            wheel_state.append(wheel[0][0])
        return wheel_state

    @property
    def prev_wheel_state(self):
        wheel_state = []
        for wheel in self.__wheels:
            wheel_state.append(wheel[-1][0])
        return wheel_state

    @staticmethod
    def __step(wheel: List):
        return wheel[1:] + wheel[:1]

    def encode(self, text: str, progress_bar):
        # 将英文字母都处理成大写
        # TODO: 使用自由字符集后，免除此项
        text = text.upper()
        result = []
        char_num = 0
        # 初始化步进状态
        across_carry_point = [False for _ in range(len(self.__wheels))]
        for i in range(len(self.__wheels)):
            if self.current_wheel_state[i] == self.__wheels_carry[i]:
                across_carry_point[i] = True
        for char in text:
            char_num += 1
            progress_bar.progress(char_num / len(text))
            # 不处理空白字符
            if re.match(r'\s', char):
                result.append(char)
                continue
            # 检查是否在字符集内：
            assert ord(char) in CHARS    # TODO: 使用自由字符集后，更改此处

            # 1. 连接板
            if self.__single_table:
                char_new = self.__single_table[char]
            else:
                char_new = char
            # 2. 整理输入
            output = copy(char_new)
            output = ord(output)    # TODO: 使用自由字符集后，更改此处
            p = list(zip(*self.__wheels[0]))[0].index(output)
            # 3. 正序过轮组
            for wheel in self.__wheels:
                output = wheel[p][1]
                p = list(zip(*wheel))[0].index(output)
            # 4. 反射板
            output = self.__reflector[p][1]
            p = list(zip(*self.__reflector))[0].index(output)
            # 5. 逆序过轮组
            for wheel in self.__wheels[::-1]:
                output = wheel[p][0]
                p = list(zip(*wheel))[1].index(output)
            # 6. 再过连接板，获得最终结果
            output = self.__wheels[0][p][0]
            output = chr(output)
            if self.__single_table:
                output = self.__single_table[output]
            result.append(output)

            # 步进，模拟真机的双重步进模式
            self.__wheels[0] = self.__step(self.__wheels[0])
            for i in range(1, len(self.__wheels) + 1):
                if self.current_wheel_state[i - 1] == self.__wheels_carry[i - 1]:
                    across_carry_point[i - 1] = True
            for i in range(1, len(self.__wheels)):
                if self.prev_wheel_state[i - 1] == self.__wheels_carry[i - 1] and across_carry_point[i - 1]:
                    self.__wheels[i] = self.__step(self.__wheels[i])
                    across_carry_point[i - 1] = False

        return ''.join(result) 
